from keras.datasets import cifar10
import numpy as np
from matplotlib_function.index import plot_image
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten, Conv2D, MaxPooling2D, ZeroPadding2D
from matplotlib_function.index import show_train_history, plot_images_labels_prediction

# 解决显存不足问题
import tensorflow as tf
from keras import backend as K

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
K.set_session(sess)

np.random.seed(10)

(x_img_train, y_label_train), (x_img_test, y_label_test) = cifar10.load_data()
# plot_image(x_img_train[0])

x_img_train_normalize = x_img_train.astype('float32') / 255.0
x_img_test_normalize = x_img_test.astype('float32') / 255.0

y_TrainOneHot = np_utils.to_categorical(y_label_train)
y_TestOneHot = np_utils.to_categorical(y_label_test)

model = Sequential()
model.add(Conv2D(filters=32, kernel_size=(3, 3), input_shape=(32, 32, 3), activation='relu', padding='same'))
model.add(Dropout(0.3))
model.add(Conv2D(filters=32, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(Dropout(0.3))
model.add(Conv2D(filters=64, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=128, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(Dropout(0.3))
model.add(Conv2D(filters=128, kernel_size=(3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dropout(0.3))
model.add(Dense(1024, activation='relu'))
model.add(Dropout(0.3))
model.add(Dense(10, activation='softmax'))
print(model.summary())

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
train_history = model.fit(x=x_img_train_normalize, y=y_TrainOneHot, validation_split=0.2, epochs=50, batch_size=300,
                          verbose=1)
show_train_history(train_history, 'acc', 'val_acc')

scores = model.evaluate(x_img_test_normalize, y_TestOneHot, verbose=1)
print("准确率：" + str(scores[1]))
prediction = model.predict_classes(x_img_test_normalize)
plot_images_labels_prediction(x_img_test, y_label_test, prediction, 0, 10)
